import logging

import mindspore
import mindspore.nn as nn
from mindspore.train.callback import Callback


def get_accuracy(network, datasets):
    metric = nn.Accuracy('classification')
    for x, y in datasets:
        y_hat = network(x)
        metric.update(y_hat, y)
    return metric.eval()


def get_loss(network, loss_f, dataset):
    loss = 0
    for x, y in dataset:
        y_hat = network(x)
        loss += loss_f(y_hat, y)
    return loss


class AccuracyMonitor(Callback):
    def __init__(self, save_url, valid_data, export):
        self.save_file_pth = save_url + 'auto-save.ckpt'
        self.best_file_pth = save_url + 'best.ckpt'
        self.local_best_pt = './best.ckpt'

        self.valid_data = valid_data
        self.export = export

        self.max_valid_acc = 0

    def epoch_end(self, run_context):
        """Called after each epoch finished."""

        callback_params = run_context.original_args()

        cur_epoch_num = callback_params.cur_epoch_num
        epoch_num = callback_params.epoch_num

        network = callback_params.network

        train_data = callback_params.train_dataset
        train_accu = get_accuracy(network, train_data)
        valid_accu = get_accuracy(network, self.valid_data)

        loss = get_loss(network, callback_params.loss_fn, train_data)
        loss /= callback_params.batch_num

        if valid_accu > self.max_valid_acc:
            self.max_valid_acc = valid_accu
            mindspore.save_checkpoint(network, self.local_best_pt)
            mindspore.save_checkpoint(network, self.best_file_pth)
            logging.info('Save the best state to {}'.format(self.local_best_pt))
        elif not cur_epoch_num % 10 and cur_epoch_num:
            mindspore.save_checkpoint(network, self.save_file_pth)
            logging.debug('Auto-save state to {}'.format(self.save_file_pth))
        print('epoch:[{}/{}] Loss:{} Train Accuracy:{} Valid Accuracy:{}'.format(cur_epoch_num, epoch_num,
                                                                                        loss, train_accu, valid_accu))
